package com.lvyh.lightframe.multidatasource.datasource;

import com.alibaba.druid.pool.DruidDataSource;
import com.lvyh.lightframe.multidatasource.common.MultiDataSourceHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.context.properties.bind.Bindable;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * Dynamic data source registration
 *
 * @author lvyh 2021/03/20.
 */
public class MultiDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {
    private static final Logger logger = LoggerFactory.getLogger(MultiDataSourceRegister.class);

    /**
     * Configuration context is the tool to get configuration files
     */
    private Environment evn;

    /**
     * Store registered default data source
     */
    private DataSource defaultDatasource = null;
    /**
     * Store registered other data sources
     */
    private Map<String, DataSource> otherDataSources = new HashMap<>();

    /**
     * Parameter binding tool
     */
    private Binder binder;

    /**
     * Custom registration bean
     */
    @Override
    public void registerBeanDefinitions(AnnotationMetadata annotationMetadata, BeanDefinitionRegistry beanDefinitionRegistry) {
        DataSource defaultDatasource = null;

        List<Map> multiDatasourceConfigList = binder.bind("multi.datasource", Bindable.listOf(Map.class)).get();
        if (!CollectionUtils.isEmpty(multiDatasourceConfigList)) {
            for (Map config : multiDatasourceConfigList) {
                DataSource tempDatasource = createDataSource(config);

                String datasourceName = config.get("name").toString();
                if (Objects.equals(datasourceName, "master")) {
                    defaultDatasource = tempDatasource;
                } else {
                    otherDataSources.put(datasourceName, tempDatasource);
                }
                // Data source context, which is used to manage data sources and record registered data source names.
                MultiDataSourceHolder.getMultiDataSources().add(datasourceName);
                logger.info("Register data source successfully:{}", datasourceName);
            }
        }


        GenericBeanDefinition define = new GenericBeanDefinition();
        // Set bean type
        define.setBeanClass(DynamicRoutingDataSource.class);

        // Parameters to be injected
        MutablePropertyValues mpv = define.getPropertyValues();

        // Add the default data source to avoid the situation that the key does not exist and no data source is available.
        mpv.add("defaultTargetDataSource", defaultDatasource);

        // Add another data source
        mpv.add("targetDataSources", otherDataSources);

        // Register the bean as a datasource without using the datasource generated automatically by springboot
        beanDefinitionRegistry.registerBeanDefinition("datasource", define);

        logger.info("Successfully registered data source, total number:{}", otherDataSources.keySet().size() + 1);
    }


    private DataSource createDataSource(Map properties) {
        DruidDataSource dataSource = new DruidDataSource();

        String initialSize = (String) properties.get("initialSize");
        if (!StringUtils.isEmpty(initialSize)) {
            dataSource.setInitialSize(Integer.parseInt(initialSize));
        }
        String maxActive = (String) properties.get("maxActive");
        if (!StringUtils.isEmpty(maxActive)) {
            dataSource.setMaxActive(Integer.parseInt(maxActive));
        }
        String maxWait = (String) properties.get("maxWait");
        if (!StringUtils.isEmpty(maxWait)) {
            dataSource.setMaxWait(Integer.parseInt(maxWait));
        }
        String name = (String) properties.get("name");
        if (!StringUtils.isEmpty(name)) {
            dataSource.setName(name);
        }
        String url = (String) properties.get("url");
        if (!StringUtils.isEmpty(url)) {
            dataSource.setUrl(url);
        }
        String username = (String) properties.get("username");
        if (!StringUtils.isEmpty(username)) {
            dataSource.setUsername(username);
        }
        String password = (String) properties.get("password");
        if (!StringUtils.isEmpty(password)) {
            dataSource.setPassword(password);
        }
        dataSource.setTestOnBorrow(true);
        dataSource.setValidationQuery("SELECT 1");
        dataSource.setDriverClassName("com.mysql.cj.jdbc.Driver");
        return dataSource;
    }

    @Override
    public void setEnvironment(Environment environment) {

        logger.info("Start registering data sources.");
        this.evn = environment;

        // Binding Configurator
        binder = Binder.get(evn);
    }
}
